import os
import sys

import torch
import torch.nn.functional
import torch.nn.modules

import numpy

from Module.xRayDetection.xray.logs import Logger
# from cv2 import Mat

class YOLOV5:

    def __init__(self, weights='', device='cpu', logger: Logger = None, git_dir='', augment=False):

        sys.path.append(git_dir if git_dir else os.path.join(os.path.dirname(__file__), 'downloads'))
        sys.path.append(os.path.join(git_dir, 'yolov5'))
        from yolov5.models import common
        from yolov5.utils import general, torch_utils, augmentations

        self.letterbox = augmentations.letterbox
        self.scale_boxes = general.scale_boxes
        self.non_max_suppression = general.non_max_suppression
        self.classify_transforms = augmentations.classify_transforms
        self.augment = augment

        try:
            self.model = common.DetectMultiBackend(weights=[w.strip() for w in weights.split(',')], device=torch_utils.select_device(device))
        except Exception as e:
            logger.log(logger.ERROR, e, exception=True)
        
        sys.path.remove(git_dir if git_dir else os.path.join(os.path.dirname(__file__), 'downloads'))
        sys.path.remove(os.path.join(git_dir, 'yolov5'))

    def to_tensor(self, img) -> torch.Tensor:
        return img

    def to_numpy(self, n) -> numpy.ndarray:
        return n

    def detect(self, image, conf_thres=0.5, box=[], iou_thres=0.45, classes=[0], x_offset=10, y_offset=30, imsz=640):

        if image is None:
            return [], []

        if len(box):
            x1 = max(int(box[0][0]) - x_offset, 0)
            x2 = min(int(box[0][2]) + x_offset, image.shape[1])
            width = x2 - x1

            y1 = max(int(box[0][1]) - y_offset, 0)
            y2 = min(y1 + width, image.shape[0])
            if y2 - y1 != width:
                y1 = image.shape[0] - width
            image = image[y1:y2, x1:x2, ...]
        
        
        img = self.to_tensor(self.letterbox(image, imsz, stride=32, auto=True)[0])

        im = img.transpose((2, 0, 1))[::-1]
        im = numpy.ascontiguousarray(im)

        im = self.to_tensor(torch.from_numpy(im).to(self.model.device))
        im = im.half() if hasattr(self.model, 'fp16') and self.model.fp16 else im.float()
        im /= 255
        if len(im.shape) == 3:
            im = im[None]

        pred = self.model(im, augment=self.augment)
        pred = self.non_max_suppression(pred, conf_thres, iou_thres)
        for det in pred:
            if len(box) == 0:
                det[:, :4] = self.to_tensor(self.scale_boxes(im.shape[2:], det[:, :4], image.shape)).round()

            det = [d for d in self.to_tensor(det).to(torch.device('cpu')).numpy()]
            if len(classes):
                det = [d for d in det if d[5] in classes]
        return img, det

    def classify(self, im, imsz=128):
        import cv2

        im = self.to_tensor(self.classify_transforms(imsz)(cv2.cvtColor(im, cv2.COLOR_BGR2RGB)))
        im = self.to_tensor(im.unsqueeze(0).to(self.model.device))
        im = im.half() if hasattr(self.model, 'fp16') and self.model.fp16 else im.float()

        results = self.model(im)

        pred = torch.nn.functional.softmax(results, dim=1)  # probabilities
        pred = self.to_numpy(pred.to(torch.device('cpu')).numpy()).tolist()[0]
        return pred

if __name__ == '__main__':
    downloads_dir = os.path.join(os.path.dirname(__file__), 'downloads')
    sys.path = [os.path.join(downloads_dir, 'yolov5')] + sys.path
    
    class Log:
        FATAL = ERROR = WARN = INFO = DEBUG = NOTSET = 0
        def log(*args, **kwargs): kwargs

    import glob
    weights = glob.glob(os.path.join(os.path.join(downloads_dir, 'models', '**', '**.pt')), recursive=True)
    model = YOLOV5(weights=weights[0], device='0', logger=Log(), yolov5_dir='', augment=False)

    import cv2
    for path in glob.glob(r''):
        im = cv2.imread(path)
        _, det = model.detect(im, 0.1)
        det
